#include "gtd_test.hpp"

int main(int argc, char* argv[])
{
    try
    {
        const std::vector<std::string> args(argv, argv + argc);
        GtdTest test(args);
        test.run();
    }
    catch(const std::exception& ex)
    {
        std::cerr << ex.what() << std::endl;
        return -1;
    }
    catch(...)
    {
        std::cerr << "Unknown error occured." << std::endl;
        return -1;
    }

    return 0;
}

static void badUsage()
{
    throw std::invalid_argument("usage: [headless] (capture | generate | <video_file>)");
}

GtdTest::GtdTest(const std::vector<std::string>& args):
isRunning(false),
graphics(true),
verticalLayout(false)
{
    // init SDL
    SDL_ERROR_CHECK(SDL_Init(0));

    // parse flag arguments
    if(args.size() < 2)
    {
        badUsage();
    }
    for(uint i = 1; i < args.size() - 1; ++i)
    {
        if(args[i] == "headless")
        {
            graphics = false;
        }
        else
        {
            badUsage();
        }
    }

    // enable CUDA/OpenGL interop
    if(graphics)
    {   
        CUDA_ERROR_CHECK(cudaGLSetGLDevice(0));
    }

    // parse arguments to init video source
    if(args.back() == "capture")
    {
        throw std::invalid_argument("Capture mode not implemented.");
    }
    else if(args.back() == "generate")
    {
        if(graphics == false)
        {
            throw std::invalid_argument("Generate and headless mode are not compatible.");
        }
        videoSource.reset(new gtd::GenVideo());
    }
    else if(args.back().find(".avi") != std::string::npos)
    {
        videoSource.reset(new gtd::AviFile(args.back()));
    }
    else
    {
        throw std::invalid_argument("Could not determine input video file format from extension.");
    }

    if(graphics)
    {
        // determine window layout
        trackerRegion.x = 0;
        trackerRegion.y = 0;
        trackerRegion.w = videoSource->getWidth();
        trackerRegion.h = videoSource->getHeight();
        verticalLayout = trackerRegion.w > trackerRegion.h;
        if(verticalLayout)
        {
            meshgenRegion.x = 0;
            meshgenRegion.y = trackerRegion.h;
        }
        else
        {
            meshgenRegion.x = trackerRegion.w;
            meshgenRegion.y = 0;
        }
        meshgenRegion.w = trackerRegion.w;
        meshgenRegion.h = trackerRegion.h;

        // init window
        SDL_ERROR_CHECK(SDL_InitSubSystem(SDL_INIT_VIDEO));
        SDL_Surface* icon = SDL_LoadBMP("data/icon1.bmp");
        SDL_ERROR_CHECK(icon == NULL);
        SDL_ERROR_CHECK(SDL_SetColorKey(icon, SDL_SRCCOLORKEY, 255));
        SDL_WM_SetIcon(icon, NULL);
        SDL_FreeSurface(icon);
        SDL_WM_SetCaption("Guess the Depth", NULL);
        
        // init video
        SDL_ERROR_CHECK(SDL_GL_SetAttribute(SDL_GL_RED_SIZE, 8));
        SDL_ERROR_CHECK(SDL_GL_SetAttribute(SDL_GL_BLUE_SIZE, 8));
        SDL_ERROR_CHECK(SDL_GL_SetAttribute(SDL_GL_GREEN_SIZE, 8));
        SDL_ERROR_CHECK(SDL_GL_SetAttribute(SDL_GL_ALPHA_SIZE, 8));
        SDL_ERROR_CHECK(SDL_GL_SetAttribute(SDL_GL_DOUBLEBUFFER, 1));                    
        SDL_Surface* screen = SDL_SetVideoMode(meshgenRegion.x + meshgenRegion.w, 
            meshgenRegion.y + meshgenRegion.h, 32, SDL_OPENGL);
        SDL_ERROR_CHECK(screen == NULL);

        // init rendering
        GL_ERROR_CHECK(glEnable(GL_TEXTURE_2D));

        GL_ERROR_CHECK(glGenTextures(1, &trackerTexture));
        GL_ERROR_CHECK(glBindTexture(GL_TEXTURE_2D, trackerTexture));
        GL_ERROR_CHECK(glTexImage2D(GL_TEXTURE_2D, 0, GL_LUMINANCE32F_ARB, 
            trackerRegion.w, trackerRegion.h, 0, GL_LUMINANCE, GL_FLOAT, NULL));
        GL_ERROR_CHECK(glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_REPLACE));
        GL_ERROR_CHECK(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST));
        GL_ERROR_CHECK(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST));
        GL_ERROR_CHECK(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT));
        GL_ERROR_CHECK(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT));
        GL_ERROR_CHECK(glBindTexture(GL_TEXTURE_2D, 0));
        
        CUDA_ERROR_CHECK(cudaGraphicsGLRegisterImage(&trackerResource, trackerTexture, 
            GL_TEXTURE_2D, cudaGraphicsMapFlagsWriteDiscard));
    }

    // init processing stages
    tracker.reset(new gtd::FeatureTracker(videoSource->getWidth(), 
        videoSource->getHeight(), 700));
    cloud.reset(new gtd::PointCloud(videoSource->getWidth(), 
        videoSource->getHeight()));
    meshgen.reset(new gtd::MeshGenerator(videoSource->getWidth(), 
        videoSource->getHeight()));
}

GtdTest::~GtdTest()
{
    PROFILER_PRINT(std::cout);

    if(graphics)
    {
        CUDA_ERROR_CHECK(cudaGraphicsUnregisterResource(trackerResource));
        GL_ERROR_CHECK(glDeleteTextures(1, &trackerTexture));
    }

    // shutdown
    SDL_Quit();
}

void GtdTest::run()
{
    SDL_Event event;
    isRunning = true;
    do
    {
        if(SDL_PollEvent(&event))
        {
            switch(event.type)
            {
                case SDL_QUIT:
                {
                    isRunning = false;
                    break;
                }
                case SDL_KEYDOWN:
                {
                    if(event.key.keysym.sym == SDLK_ESCAPE)
                    {
                        isRunning = false;
                    }
                    break;
                }
            }
        }
        else 
	    {
            update();

            PROFILER_RESET();
        }
    } 
    while(isRunning);
}

void GtdTest::update()
{
    PROFILE("Frame");

    const uchar4* d_frame = videoSource->getFrame();
    if(d_frame != NULL)
    {
        tracker->update(pairs, d_frame);
        if(!pairs.empty())
        {
            cloud->update(pairs);
            meshgen->update(*cloud, d_frame);
        }
    }
    else if(!graphics)
    {
        isRunning = false;
    }

    if(graphics)
    {
        render(d_frame != NULL);
    }
}

void GtdTest::render(bool newFrame)
{
    PROFILE("Rendering");

    GL_ERROR_CHECK(glClear(GL_COLOR_BUFFER_BIT));

    if(newFrame)
    {
        const std::vector<float*>& pyramid = tracker->getPyramid(gtd::FeatureTracker::INTENSITY);
        const std::vector<uint2>& dims = tracker->getDims();

        CUDA_ERROR_CHECK(cudaGraphicsMapResources(1, &trackerResource));
        cudaArray* trackerArray;
        CUDA_ERROR_CHECK(cudaGraphicsSubResourceGetMappedArray(&trackerArray, trackerResource, 0, 0));
        CUDA_ERROR_CHECK(cudaMemcpyToArray(trackerArray, 0, 0, pyramid[0], 
            dims[0].x * dims[0].y * sizeof(float), cudaMemcpyDeviceToDevice));
        CUDA_ERROR_CHECK(cudaGraphicsUnmapResources(1, &trackerResource));

        glBindTexture(GL_TEXTURE_2D, trackerTexture);
        glBegin(GL_QUADS);
        if(verticalLayout)
        {
            glTexCoord2f(0, 0);
            glVertex3f(-1, 0, 0);
            glTexCoord2f(1, 0);
            glVertex3f(1, 0, 0);
            glTexCoord2f(1, 1);
            glVertex3f(1, 1, 0);
            glTexCoord2f(0, 1);
            glVertex3f(-1, 1, 0);
        }
        else
        {
            glTexCoord2f(0, 0);
            glVertex3f(-1, -1, 0);
            glTexCoord2f(1, 0);
            glVertex3f(0, -1, 0);
            glTexCoord2f(1, 1);
            glVertex3f(0, 1, 0);
            glTexCoord2f(0, 1);
            glVertex3f(-1, 1, 0);
        }
        glEnd();
        glBindTexture(GL_TEXTURE_2D, 0);

        if(!pairs.empty())
        {
            glPointSize(3.0f);
            glColor4f(1, 0, 0, 1);
            glBegin(GL_POINTS);
            if(verticalLayout)
            {
                for(uint i = 0; i < pairs.size(); ++i)
                {
                    glVertex3f(2.0f * pairs[i].newPos.x / trackerRegion.w - 1.0f,
                               1.0f * pairs[i].newPos.y / trackerRegion.h - 0.0f, 0);
                }
            }
            else
            {
                for(uint i = 0; i < pairs.size(); ++i)
                {
                    glVertex3f(1.0f * pairs[i].newPos.x / trackerRegion.w - 1.0f,
                               2.0f * pairs[i].newPos.y / trackerRegion.h - 1.0f, 0);
                }
            }
            glEnd();
        }
    }

    GL_ERROR_CHECK(glFinish());
    SDL_GL_SwapBuffers();
}
